package tools

import (
	"github.com/ethereum/go-ethereum/ethclient"
	"sync"
)

type ConnectionPool struct {
	mu      sync.Mutex
	clients map[string][]*ethclient.Client
	maxSize int
}

func NewConnectionPool(size int) *ConnectionPool {
	return &ConnectionPool{
		clients: make(map[string][]*ethclient.Client),
		maxSize: size,
	}
}

func (p *ConnectionPool) Get(rpcURL string) *ethclient.Client {
	p.mu.Lock()
	defer p.mu.Unlock()

	if clients, ok := p.clients[rpcURL]; ok && len(clients) > 0 {
		client := clients[0]
		p.clients[rpcURL] = clients[1:] // 移除并返回第一个客户端
		return client
	}
	return nil // 如果没有可用连接，返回 nil
}

func (p *ConnectionPool) Put(rpcURL string, client *ethclient.Client) {
	p.mu.Lock()
	defer p.mu.Unlock()

	if clients, ok := p.clients[rpcURL]; ok && len(clients) < p.maxSize {
		p.clients[rpcURL] = append(clients, client) // 归还连接
	} else {
		client.Close() // 超过最大连接数，关闭连接
	}
}
